#!/usr/bin/env python3
import subprocess
import shutil
import os
import sys

# --- Configuration ---
CHISEL_CMD = ["sudo", "chisel", "cut", "--release", "./", "--arch", "x86_64"]
ROOTFS_DIRS = []

# --- Helper Functions ---
def run_command(cmd, check=True):
    """Runs a command, captures output, and optionally checks for errors."""
    return subprocess.run(cmd, capture_output=True, text=True, check=check)

def create_rootfs(name, slices):
    """Creates a chiseled rootfs."""
    global ROOTFS_DIRS
    rootfs_path = f"./test-{name}-rootfs"
    ROOTFS_DIRS.append(rootfs_path)
    
    run_command(CHISEL_CMD + ["--root", rootfs_path] + slices)
    # The find command is essential for making binaries and libs executable
    run_command(["sudo", "find", rootfs_path, "-type", "f", "-exec", "chmod", "a+x", "{}", "+"])
    return rootfs_path

def chroot_run(rootfs, cmd):
    """Runs a command inside a chroot environment."""
    return run_command(["sudo", "chroot", rootfs] + cmd, check=False)

def cleanup():
    """Removes all created rootfs directories."""
    for d in ROOTFS_DIRS:
        run_command(["sudo", "rm", "-rf", d])

# --- Test Definitions ---
def test_python_core():
    rootfs = create_rootfs("core", ["python3_core"])
    assert chroot_run(rootfs, ["python3", "--version"]).returncode == 0
    result = chroot_run(rootfs, ["python3", "-c", "import unittest"])
    assert result.returncode != 0 and "ModuleNotFoundError" in result.stderr, "Unittest should not be in core slice, but was found or failed for another reason."

def test_python_standard():
    rootfs = create_rootfs("standard", ["python3_standard"])
    assert chroot_run(rootfs, ["python3", "-c", "import json; import sqlite3; import unittest"]).returncode == 0

def test_python_test():
    rootfs = create_rootfs("test-unittest", ["python3_standard", "python3-test_unittest"])
    assert chroot_run(rootfs, ["python3", "-m", "unittest", "discover", "-s", "/usr/lib64/python3.11/unittest/test", "-p", "test_case.py"]).returncode == 0

def test_python_tkinter():
    # Test base libs
    rootfs_base = create_rootfs("tkinter-base", ["python3_core", "python3-tkinter_libs"])
    assert chroot_run(rootfs_base, ["python3", "-c", "import tkinter"]).returncode == 0
    result = chroot_run(rootfs_base, ["python3", "-c", "import turtle"])
    assert result.returncode != 0 and "ModuleNotFoundError" in result.stderr, "Turtle should not be in base tkinter slice, but was found or failed for another reason."
    
    # Test turtle module
    rootfs_turtle = create_rootfs("tkinter-turtle", ["python3_core", "python3-tkinter_turtle"])
    assert chroot_run(rootfs_turtle, ["python3", "-c", "import turtle"]).returncode == 0

def test_python_idle():
    rootfs = create_rootfs("idle", ["python3_standard", "python3-idle_bins"])
    assert chroot_run(rootfs, ["python3", "-c", "import idlelib"]).returncode == 0

def test_python_setuptools():
    rootfs = create_rootfs("setuptools", ["python3_standard", "python3-setuptools_libs"])
    assert chroot_run(rootfs, ["python3", "-c", "import setuptools; import pkg_resources"]).returncode == 0

def test_python_wheel():
    rootfs = create_rootfs("wheel", ["python3_standard", "python3-wheel_bins"])
    assert chroot_run(rootfs, ["wheel", "version"]).returncode == 0

def test_python_pip():
    rootfs = create_rootfs("packaging", ["python3_core", "python3-pip_bins", "python3-wheel_bins"])
    assert chroot_run(rootfs, ["pip3", "--version"]).returncode == 0
    
    # Verify bash-completion is not present
    res = chroot_run(rootfs, ["python3", "-c", "import os; assert not os.path.exists('/usr/share/bash-completion/completions/pip')"])
    assert res.returncode == 0
    
def test_python_devel():
    rootfs = create_rootfs("devel", ["python3_core", "python3-devel_headers"])
    res = chroot_run(rootfs, ["python3", "-c", "import os; assert os.path.exists('/usr/include/python3.11/Python.h')"])
    assert res.returncode == 0

def main():
    tests_to_run = [
        test_python_core,
        test_python_standard,
        test_python_test,
        test_python_tkinter,
        test_python_idle,
        test_python_setuptools,
        test_python_wheel,
        test_python_pip,
        test_python_devel,
    ]
    
    for i, test_func in enumerate(tests_to_run, 1):
        try:
            test_func()
        except Exception as e:
            print(f"\n❌ Test {i}/{len(tests_to_run)} ({test_func.__name__}) FAILED", file=sys.stderr)
            print(f"   Reason: {e}", file=sys.stderr)
            sys.exit(1)
            
    print("\n✅ All Python slice verification tests passed successfully!")

if __name__ == "__main__":
    try:
        main()
    finally:
        cleanup() 